import dask.dataframe as dd
import pandas as pd
import numpy as np
import re
import time
import sys
import traceback
from dask.diagnostics import ProgressBar
from datetime import datetime

def process_data(input_file, output_excel, age_bins, age_labels):
    """主处理函数：使用Dask处理大数据并保存到Excel"""
    # 开始计时
    start_time = time.time()
    
    # 定义需要的列（严格按提供的映射筛选）
    required_columns = [
        'albp0289',  # 服务总时长
        'axbe0031',  # 登记时间
        'axcp0002',  # 姓名
        'axcp0004',  # 人员证件类型
        'axcp0005',  # 证件号码
        'axcp0037',  # 手机号（可选，用于详细数据）
        'axcp0039',  # 出生日期
        'axcp0077'   # 注销状态（用于过滤）
    ]
    
    # 使用Dask读取CSV（自动分块并行处理）
    print(f"\n正在读取数据: {input_file}")
    ddf = dd.read_csv(
        input_file, 
        usecols=required_columns,
        dtype={
            'albp0289': 'float64',    # 服务总时长
            'axbe0031': 'object',     # 登记时间
            'axcp0002': 'object',     # 姓名
            'axcp0004': 'object',     # 人员证件类型
            'axcp0005': 'object',     # 证件号码
            'axcp0037': 'object',     # 手机号
            'axcp0039': 'object',     # 出生日期
            'axcp0077': 'object'      # 注销状态
        }
    )
    
    # 重命名列以匹配业务逻辑
    ddf = ddf.rename(columns={
        'albp0289': '服务总时长',
        'axbe0031': '登记时间',
        'axcp0002': '姓名',
        'axcp0004': '人员证件类型',
        'axcp0005': '证件号码',
        'axcp0037': '手机号',
        'axcp0039': '出生日期',
        'axcp0077': '注销状态'
    })
    
    # 筛选掉注销状态=99的记录（提前过滤）
    total_records = len(ddf)
    ddf = ddf[ddf['注销状态'].astype(str).str.strip() != '99']
    filtered_count = total_records - len(ddf)
    print(f"已过滤掉 {filtered_count} 条 注销状态=99 的记录")
    
    # 数据清洗和预处理
    print("\n正在进行数据清洗和分类...")
    ddf['人员证件类型'] = ddf['人员证件类型'].astype(str).str.strip()
    ddf['证件号码'] = ddf['证件号码'].astype(str).str.strip().str.upper()  # 统一大写
    ddf['证件长度'] = ddf['证件号码'].str.len()
    ddf['登记时间'] = dd.to_datetime(ddf['登记时间'], errors='coerce')  # 标准化日期
    
    # 数据分类（按证件类型）
    id_card_pattern = r'^(01|1|内地|居民身份证|身份证|身份证件)$'
    other_id_pattern = r'^(02|2|香港|03|3|澳门|04|4|台湾|05|5|护照|港澳|台胞)'
    
    ddf['是否内地身份证'] = ddf['人员证件类型'].str.contains(id_card_pattern, case=False, regex=True, na=False)
    ddf['是否港澳台或护照'] = ddf['人员证件类型'].str.contains(other_id_pattern, case=False, regex=True, na=False)
    
    # 提取有效身份证数据（18位）
    valid_id_ddf = ddf[ddf['是否内地身份证'] & ddf['证件长度'].eq(18)].copy()
    
    # 向量化年龄计算（核心优化）
    print("\n正在计算年龄...")
    # 从证件号码提取出生日期（18位身份证第7-14位）
    birth_dates = valid_id_ddf['证件号码'].str.extract(r'^\d{6}(\d{4})(\d{2})(\d{2})\d{4}$')
    birth_dates.columns = ['年份', '月份', '日期']
    
    # 处理闰年2月29日（向量化）
    def adjust_leap_day(row):
        if row['月份'] == '02' and row['日期'] == '29':
            year = int(row['年份'])
            if not (year % 4 == 0 and (year % 100 != 0 or year % 400 == 0)):
                return f"{row['年份']}-02-28"
        return f"{row['年份']}-{row['月份']}-{row['日期']}"
    
    birth_dates['出生日期'] = birth_dates.apply(adjust_leap_day, axis=1, meta=('出生日期', 'object'))
    birth_dates['出生日期'] = dd.to_datetime(birth_dates['出生日期'], errors='coerce')
    
    # 计算年龄（按天差/365估算）
    today = pd.Timestamp(datetime.now().date())
    valid_id_ddf['年龄'] = (today - birth_dates['出生日期']).dt.days // 365
    
    # --------------------------
    # 修正：使用Dask支持的方式赋值年龄错误类型
    # 避免直接使用loc赋值（Dask不支持）
    # --------------------------
    # 先创建默认值列，再用where替换无效年龄的记录
    valid_id_ddf['年龄错误类型'] = '成功'  # 默认值
    # 条件：无效年龄（<0或>120岁或空值）
    invalid_ages = valid_id_ddf['年龄'].isna() | (valid_id_ddf['年龄'] < 0) | (valid_id_ddf['年龄'] > 120)
    # 使用where实现条件替换（保留True的原值，替换False的值）
    valid_id_ddf['年龄错误类型'] = valid_id_ddf['年龄错误类型'].where(
        ~invalid_ages,  # 条件：有效年龄保留'成功'
        other='无效年龄（<0或>120岁）'  # 无效年龄替换为错误类型
    )
    
    # 分离有效和无效年龄记录
    failed_age_ddf = valid_id_ddf[invalid_ages].copy()
    valid_id_ddf = valid_id_ddf[~invalid_ages].copy()
    
    print(f"成功计算年龄: {len(valid_id_ddf)}条, 失败: {len(failed_age_ddf)}条")
    
    # 其他分类数据
    invalid_id_ddf = ddf[ddf['是否内地身份证'] & ~ddf['证件长度'].eq(18)].copy()  # 非18位身份证
    other_id_ddf = ddf[ddf['是否港澳台或护照']].copy()  # 港澳台/护照
    unclassified_ddf = ddf[~(ddf['是否内地身份证'] | ddf['是否港澳台或护照'])].copy()  # 未分类
    
    # 计算年龄分组统计
    print("\n正在进行年龄分组统计...")
    # 添加年龄区间列
    valid_id_ddf['年龄区间'] = dd.cut(
        valid_id_ddf['年龄'],
        bins=age_bins,
        labels=age_labels,
        right=False,
        include_lowest=True
    )
    
    # 各年龄段计数
    age_counts = valid_id_ddf['年龄区间'].value_counts().compute().reindex(age_labels).fillna(0).astype(int)
    
    # 服务时长统计（0和>0的数量）
    if '服务总时长' in valid_id_ddf.columns:
        service_stats = valid_id_ddf.groupby('年龄区间').agg(
            服务时长为0=('服务总时长', lambda x: (x == 0).sum()),
            服务时长不为0=('服务总时长', lambda x: (x > 0).sum())
        ).compute()
    else:
        service_stats = pd.DataFrame(index=age_labels)
        service_stats['服务时长为0'] = 0
        service_stats['服务时长不为0'] = 0
    
    # 证件号重复数统计
    dup_mask = valid_id_ddf.duplicated('证件号码', keep=False)
    dup_counts = valid_id_ddf[dup_mask]['年龄区间'].value_counts().compute().reindex(age_labels).fillna(0).astype(int)
    
    # 港澳台/护照类型统计
    id_type_stats = {
        '台湾': other_id_ddf['人员证件类型'].str.contains('04|4|台湾|台胞', case=False).sum().compute(),
        '香港': other_id_ddf['人员证件类型'].str.contains('02|2|香港', case=False).sum().compute(),
        '澳门': other_id_ddf['人员证件类型'].str.contains('03|3|澳门', case=False).sum().compute(),
        '护照': other_id_ddf['人员证件类型'].str.contains('05|5|护照', case=False).sum().compute()
    }
    
    # 构建最终统计结果
    age_stats = pd.DataFrame({
        '年龄段': age_labels,
        '年龄段计数': age_counts.values,
        '有效证件用户数': age_counts.values,
        '证件号缺失数': [len(invalid_id_ddf)] + [0]*(len(age_labels)-1),
        '证件号重复数': dup_counts.values,
        '台湾': [id_type_stats['台湾']] + [0]*(len(age_labels)-1),
        '香港': [id_type_stats['香港']] + [0]*(len(age_labels)-1),
        '澳门': [id_type_stats['澳门']] + [0]*(len(age_labels)-1),
        '护照': [id_type_stats['护照']] + [0]*(len(age_labels)-1),
        '未分类记录': [len(unclassified_ddf)] + [0]*(len(age_labels)-1),
    })
    
    # 合并服务时长统计
    if not service_stats.empty:
        age_stats = pd.merge(age_stats, service_stats.reset_index().rename(columns={'index': '年龄段'}), on='年龄段', how='left')
        age_stats[['服务时长为0', '服务时长不为0']] = age_stats[['服务时长为0', '服务时长不为0']].fillna(0).astype(int)
    
    # 保存结果到Excel（精简输出）
    print(f"\n正在将结果保存到Excel: {output_excel}")
    with pd.ExcelWriter(output_excel, engine='xlsxwriter') as writer:
        # 1. 统计摘要（核心结果）
        age_stats.to_excel(writer, sheet_name='统计摘要', index=False)
        
        # 2. 分类汇总（各类记录数量）
        summary_data = {
            '分类': [
                '有效内地身份证', 
                '无效年龄记录', 
                '无效内地身份证', 
                '港澳台护照', 
                '未分类记录',
                '注销状态=99的过滤记录'  # 仅保留计数
            ],
            '数量': [
                len(valid_id_ddf), 
                len(failed_age_ddf), 
                len(invalid_id_ddf),
                len(other_id_ddf), 
                len(unclassified_ddf),
                filtered_count
            ]
        }
        pd.DataFrame(summary_data).to_excel(writer, sheet_name='分类汇总', index=False)
        
        # 3. 关键详细数据（限制前10万条，避免Excel崩溃）
        print("正在保存关键详细数据（限制前10万条）...")
        max_rows = 100000
        
        # 有效身份证数据（含姓名、登记时间等）
        if not valid_id_ddf.empty:
            valid_id_df = valid_id_ddf[['姓名', '证件号码', '人员证件类型', '年龄', '年龄区间', '登记时间', '服务总时长']].compute()
            valid_id_df.head(max_rows).to_excel(writer, sheet_name='有效身份证_前10万', index=False)
        
        # 异常年龄记录（含错误原因）
        if not failed_age_ddf.empty:
            failed_age_df = failed_age_ddf[['姓名', '证件号码', '年龄错误类型', '登记时间']].compute()
            failed_age_df.head(max_rows).to_excel(writer, sheet_name='异常年龄_前10万', index=False)
    
    # 计算总耗时
    total_time = time.time() - start_time
    print(f"\n处理完成！总耗时: {total_time:.2f} 秒")
    print(f"结果已保存到Excel: {output_excel}")
    
    # 返回关键统计结果
    return {
        '总记录数': total_records,
        '过滤记录数': filtered_count,
        '有效身份证数': len(valid_id_ddf),
        '异常年龄记录数': len(failed_age_ddf),
        '耗时(秒)': total_time
    }

if __name__ == "__main__":
    # 配置参数
    input_file = r"C:\Users\17610\Desktop\test.csv"
    output_excel = r"C:\Users\17610\Desktop\数据清洗\年龄统计结果test.xlsx"
    
    # 年龄分组配置
    age_bins = [0, 6, 10, 18, 30, 40, 50, 60, 65, 70, 80, 90, 100, 120, np.inf]
    age_labels = ['0-6岁', '6-10岁', '10-18岁', '18-30岁', 
                  '30-40岁', '40-50岁', '50-60岁', 
                  '60-65岁', '65-70岁', '70-80岁', 
                  '80-90岁', '90-100岁', '100-120岁', '120岁以上']
    
    # 执行处理
    try:
        stats = process_data(input_file, output_excel, age_bins, age_labels)
        
        # 打印处理结果摘要
        print("\n===== 处理结果摘要 =====")
        print(f"总记录数: {stats['总记录数']:,}")
        print(f"过滤的注销状态=99记录数: {stats['过滤记录数']:,}")
        print(f"有效身份证记录数: {stats['有效身份证数']:,}")
        print(f"异常年龄记录数: {stats['异常年龄记录数']:,}")
        print(f"总耗时: {stats['耗时(秒)']:.2f} 秒")
        print("=======================")
        
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        traceback.print_exc()
        sys.exit(1)